﻿using UnityEngine;
using UnityEngine.Rendering;

public class Shadows
{
    private const string bufferName = "Shadow";
    private const int maxShadowedDirectionalLightCount = 4, maxShadowedOtherLightCount = 16;
    private const int maxCascades = 4;

    private static readonly string[] shadowMaskKeywords =
    {
        "_SHADOW_MASK_ALWAYS",
        "_SHADOW_MASK_DISTANCE"
    };

    private static readonly string[] directionalFilterKeywords =
    {
        "_DIRECTIONAL_PCF3",
        "_DIRECTIONAL_PCF5",
        "_DIRECTIONAL_PCF7"
    };

    private static string[] otherFilterKeywords =
    {
        "_OTHER_PCF3",
        "_OTHER_PCF5",
        "_OTHER_PCF7",
    };

    private static readonly string[] cascadeBlendKeywords =
    {
        "_CASCADE_BLEND_SOFT",
        "_CASCADE_BLEND_DITHER"
    };

    private static readonly int dirShadowAtlasId = Shader.PropertyToID("_DirectionalShadowAtlas");

    private static readonly int dirShadowMatricesId = Shader.PropertyToID("_DirectionalShadowMatrices");

    private static readonly int otherShadowAtlasId = Shader.PropertyToID("_OtherShadowAtlas");

    private static readonly int otherShadowMatricesId = Shader.PropertyToID("_OtherShadowMatrices");

    private static readonly int otherShadowTilesId = Shader.PropertyToID("_OtherShadowTiles");

    private static readonly int cascadeCountId = Shader.PropertyToID("_CascadeCount");

    private static readonly int cascadeCullingSpheresId = Shader.PropertyToID("_CascadeCullingSpheres");

    private static readonly int cascadeDataId = Shader.PropertyToID("_CascadeData");

    private static readonly int shadowAtlastSizeId = Shader.PropertyToID("_ShadowAtlasSize");

    private static readonly int shadowDistanceFadeId = Shader.PropertyToID("_ShadowDistanceFade");
    private static readonly int shadowPancakingId = Shader.PropertyToID("_ShadowPancaking");


    private static readonly Vector4[] cascadeCullingSpheres = new Vector4[maxCascades];

    private static readonly Vector4[] cascadeData = new Vector4[maxCascades];

    private static readonly Vector4[] otherShadowTiles = new Vector4[maxShadowedOtherLightCount];

    private static readonly Matrix4x4[]
        dirShadowMatrices = new Matrix4x4[maxShadowedDirectionalLightCount * maxCascades],
        otherShadowMatrices = new Matrix4x4[maxShadowedOtherLightCount];


    private readonly CommandBuffer buffer = new CommandBuffer
    {
        name = bufferName
    };

    private readonly ShadowedDirectionalLight[] shadowedDirectionalLights =
        new ShadowedDirectionalLight[maxShadowedDirectionalLightCount];

    ShadowedOtherLight[] shadowedOtherLights = new ShadowedOtherLight[maxShadowedOtherLightCount];

    private ScriptableRenderContext context;
    private CullingResults cullingResults;
    private ShadowSettings settings;

    private int ShadowedDirectionalLightCount, ShadowedOtherLightCount;

    private bool useShadowMask;

    public void Setup(ScriptableRenderContext context, CullingResults cullingResults, ShadowSettings shadowSettings)
    {
        this.context = context;
        this.cullingResults = cullingResults;
        settings = shadowSettings;
        ShadowedDirectionalLightCount = ShadowedOtherLightCount = 0;
        useShadowMask = false;
    }

    public Vector4 ReserveDirectionalShadows(Light light, int visibleLightIndex)
    {
        if (ShadowedDirectionalLightCount < maxShadowedDirectionalLightCount &&
            light.shadows != LightShadows.None && light.shadowStrength > 0f)
        {
            float maskChannel = -1;
            var lightBaking = light.bakingOutput;
            if (lightBaking.lightmapBakeType == LightmapBakeType.Mixed &&
                lightBaking.mixedLightingMode == MixedLightingMode.Shadowmask)
            {
                useShadowMask = true;
                maskChannel = lightBaking.occlusionMaskChannel;
            }

            if (!cullingResults.GetShadowCasterBounds(visibleLightIndex, out var bounds))
                return new Vector4(-light.shadowStrength, 0f, 0f, maskChannel);

            shadowedDirectionalLights[ShadowedDirectionalLightCount] =
                new ShadowedDirectionalLight
                {
                    visibleLightIndex = visibleLightIndex,
                    slopeScaleBias = light.shadowBias,
                    nearPlaneOffset = light.shadowNearPlane
                };
            return new Vector4(light.shadowStrength,
                settings.directional.cascadeCount * ShadowedDirectionalLightCount++
                , light.shadowNormalBias, maskChannel);
        }

        return new Vector4(0f, 0f, 0f, -1f);
    }

    public Vector4 ReserveOtherShadows(Light light, int visibleLightIndex)
    {
        if (light.shadows == LightShadows.None || light.shadowStrength <= 0) return new Vector4(0f, 0f, 0f, -1f);

        var maskChannel = -1f;

        var lightBaking = light.bakingOutput;
        if (lightBaking.lightmapBakeType == LightmapBakeType.Mixed &&
            lightBaking.mixedLightingMode == MixedLightingMode.Shadowmask)
        {
            useShadowMask = true;
            maskChannel = lightBaking.occlusionMaskChannel;
        }

        var isPoint = light.type == LightType.Point;
        var newLightCount = ShadowedOtherLightCount + (isPoint ? 6 : 1);
        if (newLightCount >= maxShadowedOtherLightCount ||
            !cullingResults.GetShadowCasterBounds(visibleLightIndex, out var b))
            return new Vector4(-light.shadowStrength, 0f, 0f, maskChannel);

        shadowedOtherLights[ShadowedOtherLightCount] = new ShadowedOtherLight
        {
            visibleLightIndex = visibleLightIndex,
            slopeScaleBias = light.shadowBias,
            normalBias = light.shadowNormalBias,
            isPoint = isPoint
        };
        var data = new Vector4(
            light.shadowStrength, ShadowedOtherLightCount,
            isPoint ? 1f : 0f, maskChannel);
        ShadowedOtherLightCount = newLightCount;
        return data;
    }

    private void ExecuteBuffer()
    {
        context.ExecuteCommandBuffer(buffer);
        buffer.Clear();
    }

    private Vector4 atlasSizes;

    public void Render()
    {
        if (ShadowedDirectionalLightCount > 0) RenderDirectionalShadows();
        if (ShadowedOtherLightCount > 0) RenderOtherShadows();

        buffer.BeginSample(bufferName);
        SetKeyWords(shadowMaskKeywords, useShadowMask
            ? QualitySettings.shadowmaskMode == ShadowmaskMode.Shadowmask ? 0 : 1
            : -1);
        buffer.SetGlobalInt(cascadeCountId, ShadowedDirectionalLightCount > 0 ? settings.directional.cascadeCount : 0);
        var f = 1f - settings.directional.cascadeFade;
        buffer.SetGlobalVector(shadowDistanceFadeId,
            new Vector4(1f / settings.maxDistance, 1f / settings.distanceFade, 1f / (1f - f * f)));

        buffer.SetGlobalVector(shadowAtlastSizeId, atlasSizes);
        buffer.EndSample(bufferName);
        ExecuteBuffer();
    }

    public void Cleanup()
    {
        if (ShadowedDirectionalLightCount > 0)
        {
            buffer.ReleaseTemporaryRT(dirShadowAtlasId);
            ExecuteBuffer();
        }

        if (ShadowedOtherLightCount > 0)
        {
            buffer.ReleaseTemporaryRT(otherShadowAtlasId);
            ExecuteBuffer();
        }
    }

    private Matrix4x4 ConvertToAtlasMatrix(Matrix4x4 m, Vector2 offset, float scale)
    {
        if (SystemInfo.usesReversedZBuffer)
        {
            m.m20 = -m.m20;
            m.m21 = -m.m21;
            m.m22 = -m.m22;
            m.m23 = -m.m23;
        }

        // var scale = 1f / split;
        m.m00 = (0.5f * (m.m00 + m.m30) + offset.x * m.m30) * scale;
        m.m01 = (0.5f * (m.m01 + m.m31) + offset.x * m.m31) * scale;
        m.m02 = (0.5f * (m.m02 + m.m32) + offset.x * m.m32) * scale;
        m.m03 = (0.5f * (m.m03 + m.m33) + offset.x * m.m33) * scale;
        m.m10 = (0.5f * (m.m10 + m.m30) + offset.y * m.m30) * scale;
        m.m11 = (0.5f * (m.m11 + m.m31) + offset.y * m.m31) * scale;
        m.m12 = (0.5f * (m.m12 + m.m32) + offset.y * m.m32) * scale;
        m.m13 = (0.5f * (m.m13 + m.m33) + offset.y * m.m33) * scale;
        m.m20 = 0.5f * (m.m20 + m.m30);
        m.m21 = 0.5f * (m.m21 + m.m31);
        m.m22 = 0.5f * (m.m22 + m.m32);
        m.m23 = 0.5f * (m.m23 + m.m33);
        return m;
    }


    private void RenderDirectionalShadows()
    {
        var atlasSize = (int) settings.directional.atlasSize;
        atlasSizes.x = atlasSize;
        atlasSizes.y = 1f / atlasSize;
        buffer.GetTemporaryRT(dirShadowAtlasId, atlasSize, atlasSize,
            32, FilterMode.Bilinear, RenderTextureFormat.Shadowmap);
        buffer.SetRenderTarget(dirShadowAtlasId, RenderBufferLoadAction.DontCare, RenderBufferStoreAction.Store);
        buffer.ClearRenderTarget(true, false, Color.clear);
        buffer.SetGlobalFloat(shadowPancakingId, 1f);
        buffer.BeginSample(bufferName);
        ExecuteBuffer();
        var tiles = ShadowedDirectionalLightCount * settings.directional.cascadeCount;
        var split = tiles <= 1 ? 1 : tiles <= 4 ? 2 : 4;
        var tileSize = atlasSize / split;

        for (var i = 0; i < ShadowedDirectionalLightCount; i++) RenderDirectionalShadows(i, split, tileSize);

        // buffer.SetGlobalInt(cascadeCountId, settings.directional.cascadeCount);
        buffer.SetGlobalVectorArray(cascadeCullingSpheresId, cascadeCullingSpheres);
        buffer.SetGlobalVectorArray(cascadeDataId, cascadeData);
        buffer.SetGlobalMatrixArray(dirShadowMatricesId, dirShadowMatrices);
        // var f = 1f - settings.directional.cascadeFade;
        // buffer.SetGlobalVector(shadowDistanceFadeId,
        //     new Vector4(1 / settings.maxDistance, 1f / settings.distanceFade, 1f / (1f - f * f)));
        SetKeyWords(directionalFilterKeywords, (int) settings.directional.filter - 1);
        SetKeyWords(cascadeBlendKeywords, (int) (settings.directional.cascadeBlend - 1));
        // buffer.SetGlobalVector(shadowAtlastSizeId, new Vector4(atlasSize, 1f / atlasSize));
        buffer.EndSample(bufferName);
        ExecuteBuffer();
    }

    private void RenderOtherShadows()
    {
        var atlasSize = (int) settings.other.atlasSize;
        atlasSizes.z = atlasSize;
        atlasSizes.w = 1f / atlasSize;
        buffer.GetTemporaryRT(otherShadowAtlasId, atlasSize, atlasSize,
            32, FilterMode.Bilinear, RenderTextureFormat.Shadowmap);
        buffer.SetRenderTarget(otherShadowAtlasId, RenderBufferLoadAction.DontCare, RenderBufferStoreAction.Store);
        buffer.ClearRenderTarget(true, false, Color.clear);
        buffer.SetGlobalFloat(shadowPancakingId, 0f);
        buffer.BeginSample(bufferName);
        ExecuteBuffer();
        var tiles = ShadowedOtherLightCount;
        var split = tiles <= 1 ? 1 : tiles <= 4 ? 2 : 4;
        var tileSize = atlasSize / split;

        for (var i = 0; i < ShadowedOtherLightCount;)
        {
            if (shadowedOtherLights[i].isPoint)
            {
                RenderPointShadows(i, split, tileSize);
                i += 6;
            }
            else
            {
                RenderSpotShadows(i, split, tileSize);
                i += 1;
            }
        }

        buffer.SetGlobalMatrixArray(otherShadowMatricesId, otherShadowMatrices
        );
        buffer.SetGlobalVectorArray(otherShadowTilesId, otherShadowTiles);

        SetKeyWords(otherFilterKeywords, (int) settings.other.filter - 1);

        buffer.EndSample(bufferName);
        ExecuteBuffer();
    }


    private void SetKeyWords(string[] keywords, int enabledIndex)
    {
        for (var i = 0; i < keywords.Length; i++)
            if (enabledIndex == i)
                buffer.EnableShaderKeyword(keywords[i]);
            else
                buffer.DisableShaderKeyword(keywords[i]);
    }

    private Vector2 SetTileViewport(int index, int split, float tileSize)
    {
        var offset = new Vector2(index % split, index / split);
        buffer.SetViewport(new Rect(offset.x * tileSize, offset.y * tileSize, tileSize, tileSize));
        return offset;
    }

    private void RenderDirectionalShadows(int index, int split, int tileSize)
    {
        var light = shadowedDirectionalLights[index];
        var shadowSettings = new ShadowDrawingSettings(cullingResults, light.visibleLightIndex);


        var cascadeCount = settings.directional.cascadeCount;
        var tileOffset = index * cascadeCount;
        var ratios = settings.directional.CascadeRatios();
        var cullingFactor = Mathf.Max(0f, 0.8f - settings.directional.cascadeFade);
        float tileScale = 1f / split;
        for (var i = 0; i < cascadeCount; i++)
        {
            cullingResults.ComputeDirectionalShadowMatricesAndCullingPrimitives(
                light.visibleLightIndex, i, cascadeCount, ratios,
                tileSize, light.nearPlaneOffset, out var viewMatrix, out var projectionMatrix, out var splitData);
            splitData.shadowCascadeBlendCullingFactor = cullingFactor;
            shadowSettings.splitData = splitData;
            if (index == 0) SetCascadeData(i, splitData.cullingSphere, tileSize);

            var tileIndex = tileOffset + i;
            dirShadowMatrices[tileIndex] =
                ConvertToAtlasMatrix(projectionMatrix * viewMatrix,
                    SetTileViewport(tileIndex, split, tileSize), tileScale);
            buffer.SetViewProjectionMatrices(viewMatrix, projectionMatrix);
            buffer.SetGlobalDepthBias(0, light.slopeScaleBias);
            ExecuteBuffer();
            context.DrawShadows(ref shadowSettings);
            buffer.SetGlobalDepthBias(0f, 0f);
        }
    }

    void SetOtherTileData(int index, Vector2 offset, float scale, float bias)
    {
        var border = atlasSizes.w * 0.5f;
        var data = new Vector4(
            offset.x * scale + border,
            offset.y * scale + border,
            scale - border - border,
            bias);
        otherShadowTiles[index] = data;
    }

    void RenderSpotShadows(int index, int split, int tileSize)
    {
        var light = shadowedOtherLights[index];
        var shadowSettings = new ShadowDrawingSettings(cullingResults, light.visibleLightIndex);
        cullingResults.ComputeSpotShadowMatricesAndCullingPrimitives(
            light.visibleLightIndex, out var viewMatrix, out var projMatrix, out var shadowSplitData);
        shadowSettings.splitData = shadowSplitData;

        float texelSize = 2f / (tileSize * projMatrix.m00);
        var filterSize = texelSize * ((float) settings.other.filter + 1f);
        float bias = light.normalBias * filterSize * 1.4142136f;
        Vector2 offset = SetTileViewport(index, split, tileSize);
        var scale = 1f / split;
        SetOtherTileData(index, offset, scale, bias);
        otherShadowMatrices[index] =
            ConvertToAtlasMatrix(projMatrix * viewMatrix, offset, scale);
        buffer.SetViewProjectionMatrices(viewMatrix, projMatrix);
        buffer.SetGlobalDepthBias(0, light.slopeScaleBias);
        ExecuteBuffer();
        context.DrawShadows(ref shadowSettings);
        buffer.SetGlobalDepthBias(0f, 0f);
    }

    private void RenderPointShadows(int index, int split, int tileSize)
    {
        var light = shadowedOtherLights[index];
        var shadowSettings = new ShadowDrawingSettings(cullingResults, light.visibleLightIndex);

        var texelSize = 2f / tileSize;
        var filterSize = texelSize * ((float) settings.other.filter + 1f);
        var bias = light.normalBias * filterSize * 1.4142136f;
        var scale = 1f / split;
        var fovBias = Mathf.Atan(1 + bias + filterSize) * Mathf.Rad2Deg * 2f - 90f;
        for (int i = 0; i < 6; i++)
        {
            cullingResults.ComputePointShadowMatricesAndCullingPrimitives(
                light.visibleLightIndex, (CubemapFace) i, fovBias, out var viewMatrix, out var projMatrix,
                out var shadowSplitData);
            shadowSettings.splitData = shadowSplitData;
            viewMatrix.m11 = -viewMatrix.m11;
            viewMatrix.m12 = -viewMatrix.m12;
            viewMatrix.m13 = -viewMatrix.m13;

            var tileIndex = index + i;
            Vector2 offset = SetTileViewport(tileIndex, split, tileSize);
            SetOtherTileData(tileIndex, offset, scale, bias);
            otherShadowMatrices[tileIndex] =
                ConvertToAtlasMatrix(projMatrix * viewMatrix, offset, scale);
            buffer.SetViewProjectionMatrices(viewMatrix, projMatrix);
            buffer.SetGlobalDepthBias(0, light.slopeScaleBias);
            ExecuteBuffer();
            context.DrawShadows(ref shadowSettings);
            buffer.SetGlobalDepthBias(0f, 0f);
        }
    }

    private void SetCascadeData(int index, Vector4 cullingSphere, int tileSize)
    {
        var texelSize = 2f * cullingSphere.w / tileSize;
        var filterSize = texelSize * ((float) settings.directional.filter + 1f);
        cullingSphere.w -= filterSize;
        cullingSphere.w *= cullingSphere.w;
        cascadeCullingSpheres[index] = cullingSphere;
        cascadeData[index] = new Vector4(
            1f / cullingSphere.w, filterSize * 1.4142136f
        );
    }

    private struct ShadowedDirectionalLight
    {
        public int visibleLightIndex;
        public float slopeScaleBias;
        public float nearPlaneOffset;
    }

    private struct ShadowedOtherLight
    {
        public int visibleLightIndex;
        public float slopeScaleBias;
        public float normalBias;
        public bool isPoint;
    }
}